import os, shutil
import sys
import os.path as osp
import json
import numpy as np
import cv2
import copy
from multiprocessing import Pool
from tqdm import tqdm
from tool import filesystem, via_tool, opencv_tool # export PYTHONPATH=$PYTHONPATH:`pwd`


class RectChecker:
    def __init__(self, cv_img, ignore_rects):
        """
        ignore_rects: [[x1, y1, x2, y2], ...]
        """
        self.cv_img = cv_img
        self.ignore_rects = ignore_rects
        self.mask = np.zeros(self.cv_img.shape[:2], dtype=np.uint8)
        self.mean = np.mean(np.mean(self.cv_img, axis=0), axis=0).astype(np.uint8)
        # init draw
        for rect in ignore_rects:
            cv2.rectangle(self.mask, (rect[0], rect[1]), (rect[2], rect[3]), [255], -1)

    def iou(self, mask_rect, vihicle_rect):
        """
        rect1: [x1, y1, x2, y2]
        rect2: [x1, y1, x2, y2]
        使用 vihicle_rect 作为分母 计算交并比
        """
        y0 = np.max([mask_rect[1], vihicle_rect[1]])
        y1 = np.min([mask_rect[3], vihicle_rect[3]])

        x0 = np.max([mask_rect[0], vihicle_rect[0]])
        x1 = np.min([mask_rect[2], vihicle_rect[2]])
        if x1 > x0 and y1 > y0:
            ratio = (x1 - x0) * (y1 - y0) / ((vihicle_rect[2] - vihicle_rect[0] ) * (vihicle_rect[3] - vihicle_rect[1]))
        else:
            ratio = -1
        return ratio, [x0, y0, x1, y1]
        
    def check_and_update(self, box):
        """ 绘制mask """
        for idx, rect in enumerate(self.ignore_rects):
            ratio, cover_box = self.iou(box, rect)
            cover_box = [cover_box[0]-3, cover_box[1]-3, cover_box[2]+3, cover_box[3]+3]
            if ratio > 0.8: return False
            elif ratio > 0.:
                self.mask[cover_box[1]:cover_box[3], cover_box[0]:cover_box[2]] = 0
        return True

    def get_blur_img(self):
        """ 生成 """
        paste_img = np.zeros(self.cv_img.shape, dtype=np.uint8) + self.mean
        cv2.copyTo(paste_img, self.mask, self.cv_img)
        return self.cv_img

    def draw_mask(self, box, boxs, idx):
        # check iou
        draw = True
        for i in range(len(boxs)):
            if i == idx: continue
            ratio, cover_box = self.iou(box, boxs[i])
            if ratio > 0.:
                draw = False
                break
        if draw:
            self.cv_img[box[1]:box[3], box[0]:box[2]] = self.mean

def txt_to_via(data_dir, 
                save_dir, 
                sub_dir="ua_detrac", 
                via_name="via_region_data.json"):
    """
    ua detrac tool 
    将此数据集转换成via格式
    """
    image_dir = osp.join(data_dir, "Insight-MVT_Annotation_Train")
    label_gt= osp.join(data_dir, "train_gt.txt")
    label_ign= osp.join(data_dir, "train_ign.txt")
    assert osp.exists(image_dir)
    assert osp.exists(label_gt)
    assert osp.exists(label_ign)

    # 每一张图的数据
    data_dict = dict()
    # 每一个文件夹的过滤数据
    ignore_dict = dict()
    with open(label_gt, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            data_dict[info[0]] = info[1:]
    with open(label_ign, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            ignore_dict[info[0]] = info[1:]

    sub_names = list(os.listdir(image_dir))
    thread_count  = min(12, len(sub_names)) 
    p = Pool(thread_count)
    for sub_name in sub_names:
        p.apply_async(thread_deal, args=(sub_name, image_dir, label_gt, label_ign, data_dict, ignore_dict, sub_dir, via_name))
    # print('Waiting for all subprocesses done...')
    p.close()
    p.join()


def thread_deal(sub_name, image_dir, label_gt, label_ign, data_dict, ignore_dict, sub_dir, via_name):
    ignore_rects = ignore_dict.get(sub_name, None)
    if ignore_rects is None:
        print("label_ign is None: ", sub_name)
    else:
        ignore_rects = np.array(ignore_rects, dtype=np.float32).reshape(-1, 4).astype(np.int)

    cur_dir = osp.join(image_dir, sub_name)
    cur_save_dir = osp.join(save_dir, sub_dir, sub_name)
    os.makedirs(cur_save_dir, exist_ok=True)

    total_json = dict()
    for file_name in tqdm(os.listdir(cur_dir)):
        boxs = data_dict.get(osp.join(sub_name, file_name), [])
        if boxs:
            boxs = np.array(boxs, dtype=np.float32).astype(np.int).reshape(-1, 4)

        # 保留概率
        if np.random.randint(20) != 0: continue
        file_path = osp.join(cur_dir, file_name)
        save_path = osp.join(cur_save_dir, file_name)
        cv_img = cv2.imread(file_path)
        rc = RectChecker(copy.copy(cv_img), ignore_rects)
        
        one_json = dict()
        one_json["filename"] = file_name
        one_json["file_attributes"] = {}
        
        regions = []
        for idx, box in enumerate(boxs):
            if not rc.check_and_update(box): continue

            region = dict()
            w = int(box[2] - box[0])
            h = int(box[3] - box[1])
            if w < 68 or h < 68: 
                rc.draw_mask(box, boxs, idx)
                continue
            shape_attributes = {
                "name": "rect",
                "x": int(box[0]),
                "y": int(box[1]),
                "width": w,
                "height": h
            }
            region_attributes = { "label": "car" }
            region["shape_attributes"] = shape_attributes
            region["region_attributes"] = region_attributes
            regions.append(region)

        cv2.imwrite(save_path, rc.get_blur_img())
        one_json["regions"] = regions
        file_size = osp.getsize(save_path)
        one_json["size"] = file_size
        total_json[file_name + str(file_size)] = one_json

    with open(cur_save_dir + os.sep + via_name, "w") as wf:
        wf.write(json.dumps(total_json))

def show_img_add_rect(data_dir, sub_name):
    """
    通过连续播放图片来增加过滤框
    用于过滤行人区域，避免标记
    """
    image_dir = osp.join(data_dir, "Insight-MVT_Annotation_Train")
    label_gt= osp.join(data_dir, "train_gt.txt")
    label_ign= osp.join(data_dir, "train_ign.txt")
    assert osp.exists(image_dir)
    assert osp.exists(label_gt)
    assert osp.exists(label_ign)

    # 每一张图的数据
    data_dict = dict()
    # 每一个文件夹的过滤数据
    ignore_dict = dict()
    with open(label_gt, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            data_dict[info[0]] = np.array(info[1:], dtype=np.float32).reshape(-1, 4).astype(np.int)
    with open(label_ign, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            ignore_dict[info[0]] = np.array(info[1:], dtype=np.float32).reshape(-1, 4).astype(np.int)
        

    sub_dir = osp.join(image_dir, sub_name)
    files = [[f, int(f[3:8])] for f in os.listdir(sub_dir)]
    files = sorted(files, key=lambda x: x[1])
    for f,idx in files:
        img = cv2.imread(osp.join(sub_dir, f))
        for ign_rect in ignore_dict.get(sub_name, []): 
            cv2.rectangle(img, (ign_rect[0], ign_rect[1]), (ign_rect[2], ign_rect[3]), [0,0,255], 1)
        for rect in data_dict.get(osp.join(sub_name, f), []):
            cv2.rectangle(img, (rect[0], rect[1]), (rect[2], rect[3]), [0,255,0], 1)
        cv2.imshow("image", img)
        cv2.waitKey(10)

if __name__ == "__main__":
    # # 将数据集根据mask转换
    data_dir = "/home/xc/work/data/car/UA-DETRAC"
    save_dir = "/home/xc/work/data/car/train"
    txt_to_via(data_dir, save_dir)


    # data_dir = "/home/xc/work/data/car/UA-DETRAC"
    # sub_name = sys.argv[1]
    # show_img_add_rect(data_dir, sub_name)

